#ifndef InstEngine_H
#define InstEngine_H
#include "task_metadata.h"
#include "task.h"
#include "../abc_archive.h"
#include "../abc_query/query_builder.h"
#include "be_task.h"
#include "ae_task.h"
#include "bc_tash.h"
#include <iostream>
using namespace abc;

namespace task
{
    void loadImportDescriptors(AbckitCoreFunction *coreFunc, std::vector<std::shared_ptr<ImportDescriptor>> &importDecs)
    {
        ImportsContext context;
        AbckitCoreModule *module = abckitApi->function.functionGetModule(coreFunc);

        abckitApi->import.moduleEnumerateImports(module, &context,
                                                 [](AbckitCoreImportDescriptor *i, void *data)
                                                 {
                                                     auto *context = reinterpret_cast<ImportsContext *>(data);
                                                     auto importDescriptor = std::make_shared<ImportDescriptor>();
                                                     importDescriptor->i = i;
                                                     importDescriptor->importName = abckitApi->string.getImportDescriptorName(i);
                                                     importDescriptor->importedModule = abckitApi->import.importDescriptorGetImportedModule(i);
                                                     context->imports.push_back(importDescriptor);
                                                     return true;
                                                 });
        for (auto &import : context.imports)
        {
            importDecs.push_back(import);
        }
    };

    class TaskFactory
    {
    public:
        static std::shared_ptr<TransformTask> createTask(std::shared_ptr<task::TaskParam> taskParam,
                                                         AbckitCoreFunction *coreFunc, AbckitGraph *graph,
                                                         std::vector<std::shared_ptr<ImportDescriptor>> &importDecs)
        {
            switch (taskParam->taskType)
            {
            case TaskType::BEFORE_EXECUTION:
                return std::make_shared<BeforeExecutionTask>(taskParam, coreFunc, graph, importDecs);
            case TaskType::AFTER_EXECUTION:
                return std::make_shared<AfterExecutionTask>(taskParam, coreFunc, graph, importDecs);
            case TaskType::BEFORE_CALL:
                return std::make_shared<BeforeCallTask>(taskParam, coreFunc, graph, importDecs);
            default:
                throw std::invalid_argument("Unknown TaskType");
            }
        }
    };

    using FunctionToTaskParamsMap = std::map<
        AbcFunctionPtr,
        std::set<std::shared_ptr<TaskParam>>,
        AbcFunctionPtrComparator>;
    struct TaskContext
    {
        std::vector<std::shared_ptr<task::TransformTask>> tasks;
    };
    class InstEngine
    {
    public:
        InstEngine(AbcArchive::Ptr archive) : archive_(std::move(archive)) {};
        void registerTask(std::shared_ptr<TaskParam> task)
        {
            params_.push_back(task);
        }

        void transform()
        {
            // 构建func - tasks Map
            buildFuncTaskMap();
            // 对func进行遍历，处理插桩任务
            for (const auto &pair : funcTaskParamMap_)
            {
                auto abcFunc = pair.first;
                auto taskParams = pair.second;

                handleTransformTaskByFunction(abcFunc, taskParams);
            }
        }

        std::vector<std::shared_ptr<TransformTask>> buildAtomicTasks(std::set<std::shared_ptr<task::TaskParam>> taskParams,
                                                                     AbckitCoreFunction *coreFunc, AbckitGraph *graph)
        {
            std::vector<std::shared_ptr<ImportDescriptor>> importDecs;
            loadImportDescriptors(coreFunc, importDecs);
            std::vector<std::shared_ptr<TransformTask>> tasks;
            for (auto taskParam : taskParams)
            {
                auto task = TaskFactory::createTask(taskParam, coreFunc, graph, importDecs);
                tasks.emplace_back(task);
            }
            return tasks;
        }

        void handleTransformTaskByFunction(AbcFunctionPtr abcFunc, std::set<std::shared_ptr<task::TaskParam>> taskParams)
        {
            auto coreFunc = abcFunc->coreFunction;
            auto graph = abckitApi->graph.createGraphFromFunction(abcFunc->coreFunction);
            // 构建Task，1. 分析器，每个task对应自己的分析器，分析inst、block是否命中当前task
            // 如果命中，执行当前task的插桩任务
            // 1. task任务构建 2. 当前方法体遍历 3. 遍历过程中对inst进行分析 4. 分析通过执行插桩
            auto tasks = buildAtomicTasks(taskParams, coreFunc, graph);
            AbckitBasicBlock *startBlock = abckitApi->inst.getFunctionStartBlock(graph);
            for (auto &task : tasks)
            {
                task->execute(startBlock);
            }
            TaskContext taskContext;
            taskContext.tasks = tasks;
            abckitApi->graph.gVisitBlocksRpo(graph, &taskContext, [](AbckitBasicBlock *basicBlock, void *data)
                                             {
                auto *inst = abckitApi->inst.bbGetFirstInst(basicBlock);
                while (inst != nullptr) {
                    auto *taskContext = reinterpret_cast<TaskContext *>(data);
                    for (auto &task : taskContext->tasks)
                    {
                        task->execute(inst);
                    }
                    inst = abckitApi->inst.iGetNext(inst);
                }
                return true; });

            abckitApi->function.functionSetGraph(coreFunc, graph);
            abckitApi->graph.destroyGraph(graph);
        }

    private:
        void buildFuncTaskMap()
        {
            for (auto &taskParam : params_)
            {
                auto functions = abc_query::QueryBuilder().functionName(taskParam->nodeParam->function).findFunctions(archive_);
                for (auto &func : functions)
                {
                    funcTaskParamMap_[func].insert(taskParam);
                }
            }
        }

        FunctionToTaskParamsMap funcTaskParamMap_;
        std::vector<std::shared_ptr<TaskParam>> params_;
        AbcArchive::Ptr archive_;
    };
}

#endif // InstEngine_H